# -*- coding: utf-8 -*-
import numpy as np
import pandas as pd
import pickle
from sklearn.preprocessing import StandardScaler
from prepare_data import series_to_supervised
from solve_q2_ml import init_model_grid, select_model
from sklearn.metrics import mean_squared_error, make_scorer
from sklearn.linear_model import LinearRegression
from sklearn.linear_model import LinearRegression
from sklearn.preprocessing import StandardScaler
from sklearn.svm import SVR
from sklearn.tree import DecisionTreeRegressor
from sklearn.neighbors import KNeighborsRegressor
from sklearn.ensemble import RandomForestRegressor, AdaBoostRegressor
import matplotlib.pyplot as plt

plt.rcParams['font.sans-serif']=['SimHei']
plt.rcParams['axes.unicode_minus'] = False
font = {'family' : 'SimHei',
'weight' : 'normal',
'size'   : 10,
}

def return_score(X, y, metric, verbose=True, \
            knn_param={'n_neighbors':43}, \
            svr_param={'C': 0.01, 'kernel': 'poly', 'epsilon':0.1},\
            dtr_param={'ccp_alpha':0.1, 'max_depth':4}, \
            rf_param={'n_estimators':15},\
            ada_param={'n_estimators':95}):
                
    """根据上述最优参数，构建模型"""
    lr = LinearRegression()
    knn = KNeighborsRegressor(n_neighbors=knn_param['n_neighbors'])
    svr = SVR(C=svr_param['C'], kernel=svr_param['kernel'],
              epsilon=svr_param['epsilon'])
              
    dtr = DecisionTreeRegressor(max_depth=dtr_param['max_depth'],
                                ccp_alpha=dtr_param['ccp_alpha'])
                                
    rf = RandomForestRegressor(n_estimators=rf_param['n_estimators'],\
                                max_samples=0.67,\
                                max_features=0.33, max_depth=5)
    ada = AdaBoostRegressor(base_estimator=lr,\
                            n_estimators=ada_param['n_estimators'])

    
    if verbose:
        # 画图
        estimators = ['lr', 'svr', 'knn', 'dtr', 'rf', 'ada']
        for estimator in estimators:
            predict_estimator(locals()[estimator], estimator, metric,\
                                scaler, X, y)


def predict_estimator(estimator, estimator_name,\
                    metric, scaler, X, y):
    '''
    评估各个模型，并画出实际数据、拟合数据的效果图
    '''
    estimator.fit(X, y)
    y_predict = estimator.predict(X)
    num = range(len(y.values))
    
    # 数据逆转换 
    y = scaler.inverse_transform(y)
    y_predict = scaler.inverse_transform(y_predict)
    score = metric(y, y_predict)
    
    metrics_name = [name for name in globals() if globals()[name] is metric][0]
    
    # 训练集、测试集画图
    fig = plt.figure()
    plt.plot(num, y, label='实际数据')
    plt.plot(num, y_predict, linestyle='--',\
                label=f'{estimator_name} 拟合数据')
            
    plt.xlabel('日期序列', fontsize=10)
    plt.ylabel('事故次数', fontsize=10)
    plt.title('训练集数据')
    plt.legend(prop=font)

    fig.savefig(f"../图片/第三题{estimator_name}.png")
    plt.show()

    print(f'模型 {estimator} 的 {metrics_name} 为：\n' + \
            f': {score}')



def prepare_data(time_series, problem_class):
    '''
    数据准备，导出时序数据
    '''
    # 提取某一类别的数据为例
    time_series_class = time_series.loc[time_series['事件类别']==problem_class]
    # 以月份为单位（忽视年份），统计数据
    data = time_series_class.groupby(time_series['接警日期'].dt.strftime('%B'))\
                        .count().iloc[:, 1].values
    
    month_data = pd.DataFrame(data=data, columns=['事故次数'])
    cols = month_data.columns
    
    # 数据标准化
    scaler = StandardScaler()
    month_data_after = scaler.fit_transform(month_data)
    pickle.dump(scaler, open(r'../附件/scaler_q3.pkl', 'wb'))
    # 保存标准化后的数据：
    month_data_after = pd.DataFrame(month_data_after, columns=cols)
    month_data = series_to_supervised(month_data_after, n_in=2, \
                                    n_out=1, dropnan=True)
    month_data.to_excel(f'../附件/事故类{problem_class}的月发生次数.xlsx')
    
    return month_data, scaler

if __name__ == '__main__':
    path = r'../附件/附件2：某地消防救援出警数据.xlsx'
    time_series = pd.read_excel(path)
    # 找出某个
    # 根据月份找出事故次数
    
    problem_classes = list(np.unique(time_series['事件类别']))

    # 仅分析一个问题作为示例：
    problem_class = problem_classes[0]
    print(f'本次案例主要分析 {problem_class} 问题类别，其他类别都一个模样，大家举一反三')
    
    # 数据准备, 返回事故类型 problem_class 的
    month_data, scaler = prepare_data(time_series, problem_class)

    # 训练数据、测试数据
    X = month_data.iloc[:, :-1]
    y = month_data.iloc[:, -1]
    # 初始化网格寻优法
    models_grid = init_model_grid()
    # 均方误差
    scorer = make_scorer(mean_squared_error)
    # 筛选最合适参数
    knn_param, svr_param, dtr_param, \
                    rf_param, ada_param = select_model(X, \
                                y, scorer, models_grid, cv=10, verbose=True)
    

    # 返回个模型的评分
    return_score(X, y, mean_squared_error, \
                scaler,
                knn_param=knn_param, svr_param=svr_param, \
                dtr_param=dtr_param, rf_param=rf_param, \
                ada_param=ada_param)
    print(month_data)
